#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File    :   clear_tables.py
@Time    :   2023/03/29 23:06:25
@Author  :   youle
@Email   :   a_fly0505@163.com
@Copyright : 侵权必究
'''

from core.config import Config
import psycopg2

DELETE_MARKET = '''select tablename from pg_tables where schemaname='stock_market';'''
DELETE_DATAS = '''select tablename from pg_tables where schemaname='stock_datas';'''
DROP_TABLE = '''drop table if exists {}.{};'''


class ClearDatas:
    def __init__(self, config: Config) -> None:
        self.conn = psycopg2.connect(database=config.database,
                                     user=config.user, password=config.password,
                                     host=config.host, port=config.port)
        self.cur = self.conn.cursor()

    def _clearTables(self, exec_sql: str, schema: str):
        self.cur.execute(exec_sql)
        self.conn.commit()
        tables = self.cur.fetchall()
        for table in tables:
            delete_sql = DROP_TABLE.format(schema, table[0])
            print("===", delete_sql)
            self.cur.execute(delete_sql)
            self.conn.commit()

    def ClearMarket(self):
        DROP_Market = '''drop table if exists stock_market.edge, stock_market.node, stock_market.stock;'''
        self.cur.execute(DROP_Market)
        self.conn.commit()
        # self._clearTables(DELETE_MARKET, "stock_market")

    def ClearStockDatas(self):
        self._clearTables(DELETE_DATAS, "stock_datas")
